#! python
# -*- coding: utf-8 -*-
__author__ = "caiwanpeng"

"""
并发执行工具
"""
from gevent import monkey;monkey.patch_all()

import logging
import gevent
from gevent.queue import JoinableQueue, Queue


class BaseQueryPool(object):
    """
    利用gevent 并发执行sql, 引用http://www.jasonamyers.com/2016/gevent-postgres-sqlalchemy/,
    在django中sql并发是失败的，因为一个线程中的数据库连接只有一个；
    """

    def __init__(self, queryset_dict, pool_size=5):
        """
        :param queryset_list: (queryset, fields)
        :param pool_size:
        """
        assert isinstance(queryset_dict, dict)
        self.queryset_list = queryset_dict
        self.POOL_MAX = pool_size
        self.tasks = JoinableQueue()
        self.output_queue = Queue()

    def get_query(self, queryset, fields):
        results = queryset.values(*fields) if fields else queryset.values()
        return results

    def executor(self, number):
        while not self.tasks.empty():
            i, query_set = self.tasks.get()
            queryset, fields = query_set
            try:
                results = self.get_query(queryset, fields)
                self.output_queue.put({i: results})
            except Exception as exc_info:
                logging.error(exc_info)
                logging.error('Query failed')
            self.tasks.task_done()

    def overseer(self):
        for i, query_set in self.queryset_list.items():
            self.tasks.put((i, query_set))

    def run(self):
        self.running = []
        gevent.spawn(self.overseer).join()
        for i in range(self.POOL_MAX):
            runner = gevent.spawn(self.executor, i)
            runner.start()
            self.running.append(runner)

        self.tasks.join()
        for runner in self.running:
            runner.kill()


class BaseRequestPool(object):
    """
    http请求并发执行
    """

    def __init__(self, request_dict, pool_size=2):
        """
        :param pool_size:
        """
        assert isinstance(request_dict, dict)
        self.queryset_list = request_dict
        self.POOL_MAX = pool_size
        self.tasks = JoinableQueue()
        self.output_queue = Queue()

    def get_result(self, func, params):
        """
        @desc 获取计算结果
        :param func:
        :param params:
        :return:
        """
        return func(*params)

    def executor(self, number):
        while not self.tasks.empty():
            print(number)
            i, query_set = self.tasks.get()
            func, params = query_set
            try:
                results = self.get_result(func, params)
                self.output_queue.put({i: results})
            except Exception as exc_info:
                logging.error(exc_info)
                logging.error('request failed')
            self.tasks.task_done()

    def overseer(self):
        for i, query_set in self.queryset_list.items():
            self.tasks.put((i, query_set))

    def run(self):
        self.running = []
        gevent.spawn(self.overseer).join()
        for i in range(self.POOL_MAX):
            runner = gevent.spawn(self.executor, i)
            runner.start()
            self.running.append(runner)

        self.tasks.join()
        for runner in self.running:
            runner.kill()


class LocalConcurrentManager(object):
    pool = BaseQueryPool
    req = BaseRequestPool

    @classmethod
    def run_sql(cls, queryset_dict, pool_size=5):
        """
        @desc 并发执行sql
        :param queryset_dict:
        :param pool_size:
        :return:
        """
        manager = cls.pool(queryset_dict=queryset_dict, pool_size=pool_size)
        manager.run()
        res = {}
        while not manager.output_queue.empty():
            res.update(manager.output_queue.get())
        return {k: res.get(k, None) for k in queryset_dict.keys()}

    @classmethod
    def run_request(cls, request_dict, pool_size=2):
        assert isinstance(request_dict, dict)
        manager = cls.req(request_dict=request_dict, pool_size=pool_size)
        manager.run()
        res = {}
        while not manager.output_queue.empty():
            res.update(manager.output_queue.get())
        # 执行失败的使用None占位
        return {k: res.get(k, None) for k in request_dict.keys()}


concurrent_func = LocalConcurrentManager
